/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package cn.icanci.loopstack.damc.core.dynamic;

import static damc.objectweb.asm.ClassReader.EXPAND_FRAMES;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.*;

import cn.icanci.loopstack.damc.core.dynamic.visitor.ModifierClassVisitor;
import cn.icanci.loopstack.damc.spi.model.refresh.DynamicConditionWrapper;
import damc.objectweb.asm.ClassReader;
import damc.objectweb.asm.ClassVisitor;
import damc.objectweb.asm.ClassWriter;

/**
 * @author icanci
 * @since 1.0 Created in 2025/04/21 15:05
 */
public class DynamicAgentClassCache {

    public static final Map<Class<?>, Set<DynamicConditionWrapper>> CLASS_CONDITIONS = new HashMap<>();

    /**
     * 添加动态条件
     * 
     * @param classType       类类型
     * @param wrappers 动态条件
     */
    public static void addDynamicConditions(Class<?> classType, List<DynamicConditionWrapper> wrappers) {
        CLASS_CONDITIONS.remove(classType);
        CLASS_CONDITIONS.computeIfAbsent(classType, k -> new HashSet<>()).addAll(wrappers);
    }

    /**
     * 默认的类转换器
     * 
     * @param classfileBuffer classfileBuffer
     * @return 转换后的字节码
     */
    public static byte[] defaultTransform(byte[] classfileBuffer) {
        return classfileBuffer;
    }

    /**
     * doTransform
     * 
     * @param classBeingRedefined classBeingRedefined
     * @param classfileBuffer classfileBuffer
     * @return byte[]
     */
    public static byte[] doTransform(Class<?> classBeingRedefined, byte[] classfileBuffer) {
        try {
            if (!CLASS_CONDITIONS.containsKey(classBeingRedefined)) {
                return defaultTransform(classfileBuffer);
            }
            Set<DynamicConditionWrapper> wrappers = CLASS_CONDITIONS.get(classBeingRedefined);
            ClassReader cr = getClassReader(classBeingRedefined);
            ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);

            for (DynamicConditionWrapper wrapper : wrappers) {
                if (wrapper.getClassType() == ModifierClassVisitor.class) {
                    ClassVisitor cv = new ModifierClassVisitor(cw, classBeingRedefined, wrapper.getConditions());
                    cr.accept(cv, EXPAND_FRAMES);
                }
            }

            byte[] byteArray = cw.toByteArray();

            printNewByteArray(classBeingRedefined, byteArray);

            return byteArray;
        } catch (Throwable e) {
            return defaultTransform(classfileBuffer);
        }
    }

    private static ClassReader getClassReader(Class<?> classBeingRedefined) throws IOException {
        try {
            return new ClassReader(classBeingRedefined.getTypeName());
        } catch (Throwable e) {
            String clazzName = classBeingRedefined.getTypeName().replace('.', '/') + ".class";
            return new ClassReader(classBeingRedefined.getClassLoader().getResourceAsStream(clazzName));
        }
    }

    /**
     * printNewByteArray
     * 
     * @param classBeingRedefined classBeingRedefined
     * @param byteArray byteArray
     * @throws IOException IOException
     */
    private static void printNewByteArray(Class<?> classBeingRedefined, byte[] byteArray) throws IOException {
        // 把byteArray输出到控制台为class文件
        String currentPath = System.getProperty("user.dir");
        String pathname = currentPath + "/dynamic-codes/production/classes/";
        boolean mkdirs = new File(pathname).mkdirs();
        System.out.println("[Gen Class] mkdirs = " + mkdirs);
        String filePath = pathname + classBeingRedefined.getSimpleName() + ".class";
        writeClassFile(filePath, byteArray);
        System.out.println("[Gen Class] path = " + filePath);
    }

    /**
     * 将字节数组写入到 .class 文件
     *
     * @param filePath 文件路径
     * @param data     字节数组
     * @throws IOException 如果文件写入失败
     */
    private static void writeClassFile(String filePath, byte[] data) throws IOException {
        try (FileOutputStream fos = new FileOutputStream(filePath)) {
            fos.write(data);
        }
    }
}
